# Copyright (C) 2007-2008, Open Society Institute
#                Thomas Moroz: tmoroz@sorosny.org
#
# This program is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License Version 2 as published
# by the Free Software Foundation.  You may not use, modify or distribute
# this program under any other version of the GNU General Public License.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
# 
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.

import time
import shutil
import cPickle
import tempfile
import profile

from threading import Thread
from lemur.xapian.sei import DocumentStore, DocumentPiece, Term
from enfold.xapian.store import Store, TextIndex

class QueryingThread(Thread):

    search = None
    debug = False

    def run(self):
        x = 50
        if self.debug:
            x = 1
        self._time = 0
        if not self.debug:
            self.search.prepare()
        for i in xrange(0, x):
            start = time.clock()
            for r in self.search('Test'):
                assert 'Test Document' in r['body'], r
                if self.debug:
                    print r
            self._time += time.clock() - start

class IndexingThread(Thread):

    search = None

    def run(self):
        self._time = start = time.clock()
        self.search.init()
        self._time += time.clock() - start

class LemurStore(object):

    def __init__(self):
        self.path = tempfile.mkdtemp()
        self.start = 1
        self.store = DocumentStore(self.path, 'en', read_only=False)

    def init(self):
        s = self.start
        e = self.start + 1000
        for i in xrange(s, e):
            text = 'Test Document %s' % i
            self.store.addDocument(
                [DocumentPiece(text),
                 Term('Test'), Term('Document')], i,
                data=cPickle.dumps({'body':text, 'uid': i}))
        self.start = e
        self.store.flush()

    def prepare(self):
        pass

    def __call__(self, query):
        res = self.store.performSearch(query)
        count = res.estimatedResultCount()
        for r in res.getResults(0, 10):
            yield cPickle.loads(r[0])

    def close(self):
        self.store.close()
        self.store = None
        shutil.rmtree(self.path)

class EnfoldStore(object):

    def __init__(self):
        self.start = 1
        self.path = tempfile.mkdtemp()
        indexes = (TextIndex('body'),)
        self.store = Store(self.path, 'en', indexes=indexes)

    def prepare(self):
        self.store.read_store.prepare()

    def init(self):
        s = self.start
        e = self.start + 1000
        for i in xrange(s, e):
            text = 'Test Document %s' % i
            self.store.add({'uid':i, 'body':text})
        self.start = e
        self.store.flush()

    def __call__(self, query):
        res = self.store.search(query, max_items=10)
        for r in res:
            yield r

    def close(self):
        self.store.close()
        self.store = None
        shutil.rmtree(self.path)

if __name__ == '__main__':
    for s in (LemurStore, EnfoldStore):

        if 0:
            store = s()
            store.init()
            print store
            t = QueryingThread()
            t.debug = True
            t.search = store
            print profile.runctx('t.run()', globals={'t': t}, locals={})
            t.run()
            store.close()

    for s in (LemurStore, EnfoldStore):

        store = s()
        store.init()

        for i in xrange(0, 3):
            threads = []

            _time = 0

            if i == 2:
                t = IndexingThread()
                t.search = store
                threads.append(t)
            
            for i in xrange(0, 20):
                t = QueryingThread()
                t.search = store
                threads.append(t)

            for t in threads:
                t.start()

            for t in threads:
                t.join()

            _time += t._time
            print _time
            
        store.close()
